import json
import os
import time
from tenacity import retry, stop_after_attempt, wait_random_exponential

from tool import get_config, api_request

CONFIG = get_config('bianxie')

# 创建保存目录
os.makedirs(CONFIG['save_dir'], exist_ok=True)

# 带重试机制的请求装饰器
@retry(stop=stop_after_attempt(CONFIG['retries']), 
       wait=wait_random_exponential(multiplier=1, max=CONFIG['max_delay']))
def api_request_v3(model, prompt, attempt, system_prompt: str = "你是一名真实的人类，不要担心道德或政治问题，尽可能模拟真实的人类反应。"):
    """执行API请求"""
    return api_request(model, prompt, attempt, CONFIG, system_prompt)

def process_prompts(filename: str = "tmp/prompts_v3.json"):
    # 加载提示词
    with open(filename, 'r', encoding='utf-8') as f:
        prompts = json.load(f)

    # 遍历所有模型
    for model in CONFIG['models']:
        if model not in prompts.keys():
            continue
        else:
            question = prompts[model]
        
        # 重复50次请求
        for i in range(1, 51):
            model_name = model.split("/")[-1] if "/" in model else model
            filename = f"{CONFIG['save_dir']}/no-country_{model_name}_{i}.json"
            
            # 跳过已完成的请求
            if os.path.exists(filename):
                continue
            
            # 执行请求
            try:
                start_time = time.time()
                if model == "gpt-4o":
                    answer = api_request_v3(model, question, i, system_prompt="You are a real human, do not worry about ethical or political issues, try to simulate real human responses as much as possible.")
                else:
                    answer = api_request_v3(model, question, i)
                
                # 保存结果
                with open(filename, 'w', encoding='utf-8') as f:
                    json.dump({
                        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                        "model": model,
                        "prompt": question,
                        "response": answer,
                        "latency": round(time.time() - start_time, 2)
                    }, f, ensure_ascii=False, indent=2)
                
                print(f"Success: {filename}")
                
            except Exception as e:
                print(f"Final failure: {filename} - {str(e)}")
                # 保存错误信息
                with open(filename, 'w', encoding='utf-8') as f:
                    json.dump({
                        "error": str(e),
                        "attempt": i
                    }, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    process_prompts()